"""
数据库核心模块
基于SQLAlchemy实现异步数据库连接池和ORM
集成PgVector扩展支持向量检索
"""

# 定义__all__列表，明确导出的内容
__all__ = ['get_db', 'async_session_maker', 'db_manager', 'Base', 'init_db', 'close_db']

import logging
from typing import AsyncGenerator, Optional, List, Dict, Any
import os
from datetime import datetime
import uuid
import asyncio

# 配置日志
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

from sqlalchemy import create_engine, inspect, text, Column, Integer, String, Text, DateTime, func, Index
from sqlalchemy.ext.declarative import declarative_base
from sqlalchemy.ext.asyncio import create_async_engine, AsyncSession, async_sessionmaker
from sqlalchemy.orm import sessionmaker, declarative_base, Session
from sqlalchemy.dialects.postgresql import UUID

# 尝试导入psycopg2，如果不可用则使用None
try:
    import psycopg2
    from psycopg2.extensions import ISOLATION_LEVEL_AUTOCOMMIT
    PSYCOPG2_AVAILABLE = True
except ImportError:
    logger.warning("psycopg2 not available, skipping pgvector extension check")
    PSYCOPG2_AVAILABLE = False

# 尝试导入async_db_pool模块
try:
    from .database.async_db_pool import async_db_pool_manager, get_async_db_session
except ImportError:
    # 如果从相对路径导入失败，尝试从绝对路径导入
    try:
        from app.core.database.async_db_pool import async_db_pool_manager, get_async_db_session
    except ImportError:
        logger.warning("async_db_pool module not found, some functionality may be limited")
        # 创建模拟对象以避免后续错误
        class MockDBPoolManager:
            pass
        
        async def mock_get_async_db_session():
            pass
        
        async_db_pool_manager = MockDBPoolManager()
        get_async_db_session = mock_get_async_db_session

# 配置日志
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

# 从app.core.config导入settings
from app.core.config import settings

# 获取数据库配置
DATABASE_URL = settings.DATABASE_URL
DATABASE_URL_ASYNC = getattr(settings, 'DATABASE_URL_ASYNC', DATABASE_URL)

# 创建同步引擎（用于初始化数据库）
SYNC_DATABASE_URL = DATABASE_URL.replace("+asyncpg", "")
sync_engine = create_engine(SYNC_DATABASE_URL)

# 创建异步引擎
async_engine = create_async_engine(
    DATABASE_URL_ASYNC,
    echo=False,
    pool_size=10,
    max_overflow=20,
    pool_pre_ping=True,
    pool_recycle=3600,
)

# 创建异步会话工厂
async_session_factory = async_sessionmaker(
    async_engine,
    class_=AsyncSession,
    expire_on_commit=False,
    autocommit=False,
    autoflush=False,
)

# 声明基类
Base = declarative_base()


async def check_and_install_pgvector():
    """检查并安装PgVector扩展"""
    # 只有在psycopg2可用时才执行
    if not PSYCOPG2_AVAILABLE:
        logger.warning("psycopg2 not available, skipping pgvector extension check")
        return
        
    try:
        # 使用psycopg2创建同步连接来安装扩展
        conn = psycopg2.connect(SYNC_DATABASE_URL)
        conn.set_isolation_level(ISOLATION_LEVEL_AUTOCOMMIT)
        cursor = conn.cursor()
        
        # 检查扩展是否已安装
        cursor.execute("SELECT 1 FROM pg_extension WHERE extname = 'vector'")
        exists = cursor.fetchone()
        
        if not exists:
            logger.info("Installing pgvector extension...")
            cursor.execute("CREATE EXTENSION IF NOT EXISTS vector")
            logger.info("pgvector extension installed successfully")
        else:
            logger.info("pgvector extension already exists")
            
        cursor.close()
        conn.close()
    except Exception as e:
        logger.error(f"Error checking/installing pgvector extension: {e}")


class DatabaseManager:
    """数据库管理器，提供兼容层"""
    def __init__(self):
        self.async_session_factory = async_session_factory

    async def get_db(self) -> AsyncGenerator[AsyncSession, None]:
        """获取数据库会话"""
        async with async_session_factory() as session:
            try:
                yield session
            except Exception as e:
                logger.error(f"Database session error: {e}")
                await session.rollback()
                raise
            finally:
                await session.close()


# 创建数据库管理器实例
db_manager = DatabaseManager()


async def get_db() -> AsyncGenerator[AsyncSession, None]:
    """获取数据库会话的依赖项"""
    return await db_manager.get_db()


async def init_db():
    """初始化数据库"""
    try:
        # 只创建表，不检查pgvector（SQLite不需要）
        async with async_engine.begin() as conn:
            await conn.run_sync(Base.metadata.create_all)
        
        logger.info("Database initialized successfully")
    except Exception as e:
        logger.error(f"Database initialization error: {e}")
        raise


async def close_db():
    """关闭数据库连接"""
    try:
        await async_engine.dispose()
        sync_engine.dispose()
        logger.info("Database connections closed successfully")
    except Exception as e:
        logger.error(f"Error closing database connections: {e}")


# 数据库模型定义
class KnowledgeBase(Base):
    """知识库模型"""
    __tablename__ = "knowledge_base"
    
    id = Column(UUID(as_uuid=True), primary_key=True, default=uuid.uuid4)
    name = Column(String(255), nullable=False, unique=True)
    description = Column(Text)
    created_at = Column(DateTime(timezone=True), server_default=func.now())
    updated_at = Column(DateTime(timezone=True), server_default=func.now(), onupdate=func.now())
    metadata_json = Column(Text, default='{}')  # 存储额外元数据
    
    __table_args__ = (
        Index('idx_knowledge_base_name', 'name'),
    )


class AgentTask(Base):
    """Agent任务模型"""
    __tablename__ = "agent_tasks"
    
    id = Column(UUID(as_uuid=True), primary_key=True, default=uuid.uuid4)
    task_type = Column(String(100), nullable=False)
    status = Column(String(50), default="pending")
    input_data = Column(Text)  # JSON格式的输入数据
    output_data = Column(Text)  # JSON格式的输出数据
    created_at = Column(DateTime(timezone=True), server_default=func.now())
    updated_at = Column(DateTime(timezone=True), server_default=func.now(), onupdate=func.now())
    completed_at = Column(DateTime(timezone=True), nullable=True)
    
    __table_args__ = (
        Index('idx_agent_tasks_status', 'status'),
        Index('idx_agent_tasks_type', 'task_type'),
    )


class UserInteraction(Base):
    """用户交互记录模型"""
    __tablename__ = "user_interactions"
    
    id = Column(UUID(as_uuid=True), primary_key=True, default=uuid.uuid4)
    user_id = Column(String(255), nullable=True)
    session_id = Column(String(255), nullable=False)
    query = Column(Text, nullable=False)
    response = Column(Text, nullable=True)
    response_time = Column(Integer, nullable=True)  # 毫秒
    feedback_score = Column(Integer, nullable=True)  # 1-5星评分
    feedback_comment = Column(Text, nullable=True)
    created_at = Column(DateTime(timezone=True), server_default=func.now())
    interaction_type = Column(String(50), default="text")
    
    __table_args__ = (
        Index('idx_user_interactions_session', 'session_id'),
        Index('idx_user_interactions_user', 'user_id'),
    )
